import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from models.TransMIL.nystrom_attention import NystromAttention
from infini_transformer.infini_transformer import InfiniTransformer,YaRNEmbeddings
from torchinfo import summary


# class TransLayer(nn.Module):
#
#     def __init__(self, norm_layer=nn.LayerNorm, dim=512):
#         super().__init__()
#         self.norm = norm_layer(dim)
#         self.attn = NystromAttention(
#             dim=dim,
#             dim_head=dim // 8,
#             heads=8,
#             num_landmarks=dim // 2,  # number of landmarks
#             pinv_iterations=6,
#             # number of moore-penrose iterations for approximating pinverse. 6 was recommended by the paper
#             residual=True,
#             # whether to do an extra residual with the value or not. supposedly faster convergence if turned on
#             dropout=0.1
#         )
#
#     def forward(self, x):
#         x = x + self.attn(self.norm(x))
#
#         return x


class Translayer(nn.Module):

    def __init__(self, norm_layer=nn.LayerNorm, dim_input=512):
        super(Translayer, self).__init__()
        self.norm = norm_layer(dim_input)
        self.dim_input = dim_input
        self.dim_hidden = 2048
        self.dim_key = 64
        self.dim_value = 64
        self.num_heads = 8
        self.activation = "swiglu"
        self.segment_len = 4096
        self.update = "delta"
        self.causal = True
        self.init_state_learnable = True
        self.dropout = 0.1

        self.positional_embedder = YaRNEmbeddings(
            dim=self.dim_key,
            seq_len=self.segment_len,
            context_len=32000,
            context_len_ext=64000,
            dim_embedding_pct=0.7,
            base=10000,
            alpha=1,
            beta=32,
            length_scale=None
        )

        self.attn = InfiniTransformer(
            dim_input=self.dim_input,
            dim_hidden=self.dim_hidden,
            dim_key=self.dim_key,
            dim_value=self.dim_value,
            num_heads=self.num_heads,
            activation=self.activation,
            segment_len=self.segment_len,
            update=self.update,
            causal=self.causal,
            positional_embedder=self.positional_embedder,
            init_state_learnable=self.init_state_learnable,
            dropout=self.dropout
        )

    def forward(self, x):
        x = x + self.attn(self.norm(x))

        return x


class PPEG(nn.Module):
    def __init__(self, dim=512):
        super(PPEG, self).__init__()
        self.proj = nn.Conv2d(dim, dim, 7, 1, 7//2, groups=dim)
        self.proj1 = nn.Conv2d(dim, dim, 5, 1, 5//2, groups=dim)
        self.proj2 = nn.Conv2d(dim, dim, 3, 1, 3//2, groups=dim)

    def forward(self, x, H, W):
        B, _, C = x.shape
        cls_token, feat_token = x[:, 0], x[:, 1:]
        cnn_feat = feat_token.transpose(1, 2).view(B, C, H, W)
        x = self.proj(cnn_feat) + cnn_feat + self.proj1(cnn_feat) + self.proj2(cnn_feat)
        x = x.flatten(2).transpose(1, 2)
        x = torch.cat((cls_token.unsqueeze(1), x), dim=1)
        return x


class TransMIL_infi(nn.Module):
    def __init__(self, n_classes):
        super(TransMIL_infi, self).__init__()
        self.pos_layer = PPEG(dim=512)
        self._fc1 = nn.Sequential(nn.Linear(1024, 512), nn.ReLU())
        self.cls_token = nn.Parameter(torch.randn(1, 1, 512))
        self.n_classes = n_classes

        # 使用 InfiniTransformer 替换 TransLayer
        self.layer1 = Translayer()
        self.layer2 = Translayer()
        self.norm = nn.LayerNorm(512)
        self._fc2 = nn.Linear(512, self.n_classes)

    def forward(self, h, coords=None, **kwargs):
        h = h.unsqueeze(0).float()  # [1, n, 1024]
        # print(f"Input shape after unsqueeze: {h.shape}")

        h = self._fc1(h)  # [B, n, 512]
        # print(f"Shape after _fc1: {h.shape}")

        # ---->pad
        H = h.shape[1]
        _H, _W = int(np.ceil(np.sqrt(H))), int(np.ceil(np.sqrt(H)))
        add_length = _H * _W - H
        h = torch.cat([h, h[:, :add_length, :]], dim=1)  # [B, N, 512]
        # print(f"Shape after padding: {h.shape}")

        # ---->cls_token
        B = h.shape[0]
        cls_tokens = self.cls_token.expand(B, -1, -1).to(h.device)  # Ensure cls_tokens are on the same device as h
        h = torch.cat((cls_tokens, h), dim=1)
        # print(f"Shape after adding cls_token: {h.shape}")


        # ---->Translayer x1
        h = self.layer1(h)  # [B, N, 512]
        # print(f"Shape after layer1: {h.shape}")

        # ---->PPEG
        h = self.pos_layer(h, _H, _W)  # [B, N, 512]
        # print(f"Shape after pos_layer: {h.shape}")

        # ---->Translayer x2
        h = self.layer2(h)  # [B, N, 512]
        # print(f"Shape after layer2: {h.shape}")

        # ---->cls_token
        h = self.norm(h)[:, 0]
        # print(f"Shape after norm and selecting cls_token: {h.shape}")

        # ---->predict
        logits = self._fc2(h)  # [B, n_classes]
        Y_hat = torch.argmax(logits, dim=1)
        Y_prob = F.softmax(logits, dim=1)

        attn = []
        results_dict = {}
        return logits, Y_prob, Y_hat, attn, results_dict


if __name__ == "__main__":

    model = TransMIL_infi(n_classes=2)
    summary(model, (100, 1024), depth=3)





